import json
import logging
import re

from myapp.utils.exception import ignore_exception

DEFAULT_SCHEDULER = 'default-scheduler'


class DeviceResource:
    typ = ''

    def __init__(self, num, mem=None, node_selector=None):
        node_selector = node_selector or {}
        num = int(num)
        if mem:
            mem = int(mem)
        if num < 0:
            raise ValueError('Number of GPUs must be greater than 0')
        if mem and mem < 0:
            raise ValueError('Memory must be greater than 0')
        if node_selector is None:
            node_selector = {}

        self.num = int(num)
        self.mem = mem
        self._node_selector = node_selector
        self.validate()

    def validate(self):
        pass

    @staticmethod
    def build(typ, num, mem=None, node_selector=None):
        node_selector = node_selector or {}
        if isinstance(mem, int) and mem < 0:
            raise ValueError('Memory must be greater than 0')
        if num > 0 and typ is None:
            typ = 'nvidia'
        if num == 0:
            if typ is None or typ == "":
                # 使用CPU调度
                return DeviceResource(0)

        cls = all_gpu_resources.get(typ)
        if cls is None:
            raise ValueError(f'unknown gpu type:{typ}')
        return cls(num, mem, node_selector)

    @staticmethod
    def from_str(spec, typ=None):
        logging.info(f'[from_str] spec: {spec}; typ: {typ}')
        if spec is None:
            return DeviceResource.build(typ, 0)

        # 推理常用的规格： num
        if isinstance(spec, int):
            return DeviceResource.build(typ, spec)

        with ignore_exception:
            spec = int(spec)
            return DeviceResource.build(typ, spec)

        # json格式
        with ignore_exception:
            res = json.loads(spec)
            typ = res.get('type')
            num = res.get('num')
            mem = res.get('mem')
            node_selector = res.get('node_selector', {})
            return DeviceResource.build(typ, num, mem=mem, node_selector=node_selector)

        # 推理常用的规格： num,mem
        if ',' in spec:
            gpu_mem = spec.split(',')[1]
            if int(gpu_mem) < 100:
                gpu_mem = str(1000)
            gpu_num = spec.split(',')[0]
            gpu_num = int(gpu_num)

            return DeviceResource.build(typ, gpu_num, gpu_mem)

        # 训练常用的规格：1(nvidia)[10k] 1(nvidia)
        gpu_num = 0
        gpu_type = None
        gpu_mem = None

        with ignore_exception:
            if '(' in spec:
                gpu_type = re.findall(r'\((.+?)\)', spec)
                gpu_type = gpu_type[0] if gpu_type else None
            if '（' in spec:
                gpu_type = re.findall(r'（(.+?)）', spec)
                gpu_type = gpu_type[0] if gpu_type else None
            if '[' in spec:
                gpu_mem = re.findall(r'\[(.+?)]', spec)
                gpu_mem = gpu_mem[0] if gpu_mem else None

            if '(' in spec:
                resource_gpu = spec[0 : spec.index('(')]
            elif '（' in spec:
                resource_gpu = spec[0 : spec.index('（')]
            elif '[' in spec:
                resource_gpu = spec[0 : spec.index('[')]
            gpu_num = int(resource_gpu)

        if isinstance(gpu_mem, str):
            gpu_mem = gpu_mem.strip()
            if gpu_mem.endswith('k'):
                gpu_mem = gpu_mem[:-1]
                gpu_mem = int(gpu_mem) * 1000

            gpu_mem = int(gpu_mem)
            if gpu_mem == 0:
                gpu_mem = None

        return DeviceResource.build(gpu_type, gpu_num, gpu_mem)

    @property
    def node_selector(self):
        return {'cpu': 'true'}

    @property
    def resources_requests(self):
        return {}

    @property
    def resources_limits(self):
        return {}

    @property
    def volume_mounts(self):
        return {}

    @property
    def volumes(self):
        return {}

    @property
    def envs(self):
        return {'NVIDIA_VISIBLE_DEVICES': 'none'}


class NvidiaDeviceResource(DeviceResource):
    typ = 'nvidia'

    def validate(self):
        if self.num < 0:
            raise ValueError('Number of GPUs must be > 0')

    @property
    def node_selector(self):
        sel = {}
        for k, v in self._node_selector.items():
            sel[k] = v
        if len(self._node_selector) == 0:
            sel['card-type'] = self.typ
        sel['gpu'] = 'true'
        return sel

    @property
    def resources_requests(self):
        if self.num <= 0:
            return {}

        return {'nvidia.com/gpu': self.num}

    @property
    def resources_limits(self):
        res = {}
        if self.num > 0:
            res['nvidia.com/gpu'] = self.num
        if self.mem and self.mem > 0:
            gpumem = round(self.mem / 1000)
            res['nvidia.com/gpumem'] = int(gpumem)

        return res

    @property
    def volume_mounts(self):
        return {}

    @property
    def volumes(self):
        return {}

    @property
    def envs(self):
        if self.num == 0:
            return {'NVIDIA_VISIBLE_DEVICES': 'none'}
        return {}


class Ascend910DeviceResource(DeviceResource):
    typ = 'ascend910'

    def validate(self):
        if self.num <= 0:
            raise ValueError('Number of GPUs must be > 0')

    @property
    def node_selector(self):
        if self.num <= 0:
            return {}

        sel = {}
        for k, v in self._node_selector.items():
            sel[k] = v
        if len(self._node_selector) == 0:
            sel['card-type'] = self.typ
        sel['host-arch'] = 'huawei-arm'
        return sel

    @property
    def resources_requests(self):
        if self.num <= 0:
            return {}

        return {'huawei.com/Ascend910': self.num}

    @property
    def resources_limits(self):
        res = {}
        if self.num > 0:
            res['huawei.com/Ascend910'] = self.num

        return res

    @property
    def volume_mounts(self):
        if self.num <= 0:
            return {}
        return {
            'ascend-driver': {'mountPath': '/usr/local/Ascend/driver'},
            'ascend-add-ons': {'mountPath': '/usr/local/Ascend/add-ons'},
        }

    @property
    def volumes(self):
        if self.num <= 0:
            return {}

        return {
            'ascend-driver': {'hostPath': {'path': '/usr/local/Ascend/driver', 'type': ''}},
            'ascend-add-ons': {'hostPath': {'path': '/usr/local/Ascend/add-ons', 'type': ''}},
        }


class HygonDeviceResource(DeviceResource):
    typ = 'hygon'
    resource_key = 'hygon.com/dcu'

    def validate(self):
        if self.num <= 0:
            raise ValueError('Number of DPUs must be > 0')

    @property
    def node_selector(self):
        if self.num <= 0:
            return {}

        sel = {}
        for k, v in self._node_selector.items():
            sel[k] = v
        if len(self._node_selector) == 0:
            sel['card-type'] = self.typ
        return sel

    @property
    def resources_requests(self):
        if self.num <= 0:
            return {}

        return {self.resource_key: self.num}

    @property
    def resources_limits(self):
        res = {}
        if self.num > 0:
            res[self.resource_key] = self.num

        return res

    @property
    def volume_mounts(self):
        if self.num <= 0:
            return {}
        return {
            'hygon-plugin': {'mountPath': '/opt/hyhal'},
        }

    @property
    def volumes(self):
        if self.num <= 0:
            return {}

        return {
            'hygon-plugin': {'hostPath': {'path': '/opt/hyhal', 'type': ''}},
        }


all_gpu_resources = {
    NvidiaDeviceResource.typ: NvidiaDeviceResource,
    Ascend910DeviceResource.typ: Ascend910DeviceResource,
    HygonDeviceResource.typ: HygonDeviceResource,
}
